# knn_step1_data.py
import pandas as pd
import numpy as np
from sklearn.datasets import load_iris

print("=== 第2题步骤1: Iris数据集分析 ===")

# 方法1: 使用sklearn内置的Iris数据集（推荐）
def load_iris_from_sklearn():
    print("📊 从sklearn加载Iris数据集...")
    iris = load_iris()
    X = iris.data  # 特征
    y = iris.target  # 标签
    feature_names = iris.feature_names
    target_names = iris.target_names
    
    print("✅ 数据加载成功！")
    print(f"数据形状: {X.shape}")
    print(f"特征数: {X.shape[1]}")
    print(f"样本数: {X.shape[0]}")
    
    print("\n🔍 特征名称:")
    for i, name in enumerate(feature_names):
        print(f"  特征{i+1}: {name}")
    
    print("\n🎯 类别信息:")
    for i, name in enumerate(target_names):
        count = np.sum(y == i)
        print(f"  类别{i} ({name}): {count} 个样本")
    
    print("\n📋 数据前5行:")
    data_preview = pd.DataFrame(X, columns=feature_names)
    data_preview['target'] = y
    data_preview['species'] = [target_names[i] for i in y]
    print(data_preview.head())
    
    return X, y, feature_names, target_names

# 方法2: 从UCI网站加载（备选）
def load_iris_from_uci():
    print("\n🌐 尝试从UCI网站加载Iris数据集...")
    try:
        # UCI Iris数据集的URL
        url = "https://archive.ics.uci.edu/ml/machine-learning-databases/iris/iris.data"
        column_names = ['sepal_length', 'sepal_width', 'petal_length', 'petal_width', 'class']
        iris_data = pd.read_csv(url, header=None, names=column_names)
        
        print("✅ 从UCI加载成功！")
        return iris_data
    except:
        print("❌ 从UCI加载失败，使用sklearn数据")
        return None

# 主程序
try:
    # 使用sklearn数据（更稳定）
    X, y, feature_names, target_names = load_iris_from_sklearn()
    
    # 保存数据到本地文件
    iris_df = pd.DataFrame(X, columns=feature_names)
    iris_df['target'] = y
    iris_df['species'] = [target_names[i] for i in y]
    iris_df.to_csv('iris_data.csv', index=False)
    print(f"\n💾 数据已保存到: iris_data.csv")
    
    # 数据统计分析
    print("\n📈 数据统计分析:")
    print("各特征的统计信息:")
    stats_df = pd.DataFrame(X, columns=feature_names)
    print(stats_df.describe())
    
    # 类别分布可视化
    import matplotlib.pyplot as plt
    
    plt.figure(figsize=(10, 6))
    
    # 选择两个特征进行可视化
    colors = ['red', 'blue', 'green']
    for i in range(3):
        mask = (y == i)
        plt.scatter(X[mask, 0], X[mask, 1], 
                   c=colors[i], label=target_names[i], 
                   alpha=0.7, s=50)
    
    plt.xlabel(feature_names[0])
    plt.ylabel(feature_names[1])
    plt.title('Iris数据集可视化 (前两个特征)')
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.savefig('iris_visualization.png', dpi=300, bbox_inches='tight')
    plt.show()
    
    print("✅ 可视化图已保存为: iris_visualization.png")
    
    print(f"\n🎯 数据集总结:")
    print(f"- 总样本数: {len(X)}")
    print(f"- 特征数: {X.shape[1]}")
    print(f"- 类别数: {len(target_names)}")
    print(f"- 特征: {', '.join(feature_names)}")
    print(f"- 类别: {', '.join(target_names)}")
    
except Exception as e:
    print(f"❌ 出错: {e}")
    import traceback
    traceback.print_exc()

input("\n按 Enter 键继续下一步...")